#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@文件        :    coco2yolo.py
@说明        :    将coco数据集转化为yolo形式
@时间        :    2021/12/20 15:11:52
@作者        :    钟建平
@版本        :    1.0
'''
import  os
import argparse
import json
import shutil
from tqdm import tqdm


def  _parser_args():
    """parse input"""
    parser = argparse.ArgumentParser("convert coco to yolo")
    parser.add_argument('--coco_path', type=str, required=True, help='the path of coco')
    parser.add_argument('--yolo_path', type=str,  required=True,  help='the path ot save yolo data')
    args, _ = parser.parse_known_args()
    return args

def convert(coco_path, yolo_path):
    """
            convert coco to yolo
            coco数据集格式和yolo数据集格式
            coco
                |----annotations
                            |----instances_train2017.json
                            |----instances_val2017.json
                |----train2017
                |----val2017
            yolo
                |----images
                            |----train2017
                            |----val2017
                |----labels
                            |----train2017
                            |----val2017    
    """
    annotation_list  = ['train2017',  'val2017']
    predix = 'instances'
    if not os.path.exists(yolo_path):
        os.mkdir(yolo_path)
    if not os.path.exists(yolo_path+'images'):
        os.mkdir(yolo_path+'images')
    if not os.path.exists(yolo_path+'labels'):
        os.mkdir(yolo_path+'labels')
    for annotation_file in annotation_list:
        print(f'----------------------- start convert {annotation_file} ----------------')
        annotations = json.load(open(coco_path+'annotations/'+predix+'_'+annotation_file+'.json'))
        cat_id_map = dict()
        id_index = 0
        for cat in annotations['categories']:
            cat_id_map[cat['id']] = id_index
            id_index += 1
        if not os.path.exists(yolo_path+'images/'+annotation_file):
            os.mkdir(yolo_path+'images/'+annotation_file)
        if not os.path.exists(yolo_path+'labels/'+annotation_file):
            os.mkdir(yolo_path+'labels/'+annotation_file)
        # image_id_info, image_id_annotations
        image_id_info = dict()
        image_id_annotations = dict()
        for image in annotations['images']:
            image_id_info[image['id']] = [image['file_name'], image['width'], image['height']]
            image_id_annotations[image['id']] = []
        # image_id_annotations
        for annotation in annotations['annotations']:
            image_id_annotations[annotation['image_id']].append(annotation['bbox']+[annotation['category_id']])
        
        # 将coco标注转化为yolo格式
        for image_id in tqdm(image_id_annotations):
            objs = image_id_annotations[image_id]
            image_name = image_id_info[image_id][0]
            f = open(yolo_path+'labels/'+annotation_file+'/'+image_name.replace('.jpg', '.txt'), 'w')
            shutil.copy(coco_path+annotation_file+'/'+image_name, yolo_path+'images/'+annotation_file+'/'+image_name)
            img_w, img_h = image_id_info[image_id][1], image_id_info[image_id][2]
            for obj in objs:
                [x, y, w, h] = obj[:4]
                catgeory_id = obj[4]
                x_center = x + w / 2
                y_center = y + h/ 2
                f.write(str(cat_id_map[catgeory_id])+' '+str(x_center/img_w)+' '+str(y_center/img_h)+' '+str(w/img_w)+' '+str(h/img_h)+'\n')
            f.close()
        print(f'----------------------- convert {annotation_file}  success----------------')

if  __name__ ==  "__main__":
    """python coco2yolo.py --coco_path [coco_path] --yolo_path [yolo_path]"""
    args  = _parser_args()
    convert(args.coco_path, args.yolo_path)
